#include "color_correction/color_correction.hpp"
#include <algorithm>
#include <cmath>
#include <locale>

ColorCorrector::ColorCorrector() : isCalibrated(false) {
    // 初始化标准色卡颜色值 (24色Macbeth色卡的sRGB值)
    // 这里使用一些典型的色卡颜色作为参考
    standardColors = {
        cv::Vec3f(115, 82, 68),   // 深棕色
        cv::Vec3f(194, 150, 130), // 浅棕色
        cv::Vec3f(98, 122, 157),  // 蓝色
        cv::Vec3f(87, 108, 67),   // 绿色
        cv::Vec3f(133, 128, 177), // 浅紫色
        cv::Vec3f(103, 189, 170), // 青绿色
        cv::Vec3f(214, 126, 44),  // 橙色
        cv::Vec3f(80, 91, 166),   // 紫蓝色
        cv::Vec3f(193, 90, 99),   // 粉红色
        cv::Vec3f(94, 60, 108),   // 深紫色
        cv::Vec3f(157, 188, 64),  // 黄绿色
        cv::Vec3f(224, 163, 46),  // 黄橙色
        cv::Vec3f(56, 61, 150),   // 蓝色
        cv::Vec3f(70, 148, 73),   // 绿色
        cv::Vec3f(175, 54, 60),   // 红色
        cv::Vec3f(231, 199, 31),  // 黄色
        cv::Vec3f(187, 86, 149),  // 洋红色
        cv::Vec3f(8, 133, 161),   // 青色
        cv::Vec3f(243, 243, 242), // 白色
        cv::Vec3f(200, 200, 200), // 浅灰色
        cv::Vec3f(160, 160, 160), // 中灰色
        cv::Vec3f(122, 122, 121), // 深灰色
        cv::Vec3f(85, 85, 85),    // 更深灰色
        cv::Vec3f(52, 52, 52)     // 黑色
    };
}

cv::Mat ColorCorrector::findColorCard(const cv::Mat& image) {
    // 创建ArUco检测器
    cv::Ptr<cv::aruco::Dictionary> dictionary = cv::aruco::getPredefinedDictionary(cv::aruco::DICT_ARUCO_ORIGINAL);
    cv::Ptr<cv::aruco::DetectorParameters> parameters = cv::aruco::DetectorParameters::create();
    
    std::vector<std::vector<cv::Point2f>> corners;
    std::vector<int> ids;
    
    // 检测ArUco标记
    cv::aruco::detectMarkers(image, dictionary, corners, ids, parameters);
    
    if (ids.size() < 4) {
        std::cout << "[错误] 未能找到足够的ArUco标记 (需要4个，找到" << ids.size() << "个)" << std::endl;
        return cv::Mat();
    }
    
    // 查找特定ID的标记
    cv::Point2f topLeft, topRight, bottomRight, bottomLeft;
    bool foundAll = false;
    
    try {
        for (size_t i = 0; i < ids.size(); ++i) {
            switch (ids[i]) {
                case TOP_LEFT_ID:
                    topLeft = corners[i][0];
                    break;
                case TOP_RIGHT_ID:
                    topRight = corners[i][1];
                    break;
                case BOTTOM_RIGHT_ID:
                    bottomRight = corners[i][2];
                    break;
                case BOTTOM_LEFT_ID:
                    bottomLeft = corners[i][3];
                    break;
            }
        }
        
        // 检查是否找到所有四个标记
        foundAll = std::find(ids.begin(), ids.end(), TOP_LEFT_ID) != ids.end() &&
                   std::find(ids.begin(), ids.end(), TOP_RIGHT_ID) != ids.end() &&
                   std::find(ids.begin(), ids.end(), BOTTOM_RIGHT_ID) != ids.end() &&
                   std::find(ids.begin(), ids.end(), BOTTOM_LEFT_ID) != ids.end();
        
    } catch (const std::exception& e) {
        std::cout << "[错误] 提取ArUco标记坐标时出错: " << e.what() << std::endl;
        return cv::Mat();
    }
    
    if (!foundAll) {
        std::cout << "[错误] 未能找到所有必需的ArUco标记ID" << std::endl;
        return cv::Mat();
    }
    
    // 进行四点透视变换
    std::vector<cv::Point2f> cardCorners = {topLeft, topRight, bottomRight, bottomLeft};
    return fourPointTransform(image, cardCorners);
}

cv::Mat ColorCorrector::fourPointTransform(const cv::Mat& image, const std::vector<cv::Point2f>& corners) {
    // 计算变换后的宽度和高度
    float widthA = cv::norm(corners[2] - corners[3]);
    float widthB = cv::norm(corners[1] - corners[0]);
    int maxWidth = static_cast<int>(std::max(widthA, widthB));
    
    float heightA = cv::norm(corners[1] - corners[2]);
    float heightB = cv::norm(corners[0] - corners[3]);
    int maxHeight = static_cast<int>(std::max(heightA, heightB));
    
    // 目标矩形的四个角点
    std::vector<cv::Point2f> dst = {
        cv::Point2f(0, 0),
        cv::Point2f(maxWidth - 1, 0),
        cv::Point2f(maxWidth - 1, maxHeight - 1),
        cv::Point2f(0, maxHeight - 1)
    };
    
    // 计算透视变换矩阵
    cv::Mat transform = cv::getPerspectiveTransform(corners, dst);
    
    // 应用透视变换
    cv::Mat warped;
    cv::warpPerspective(image, warped, transform, cv::Size(maxWidth, maxHeight));
    
    return warped;
}

std::vector<cv::Vec3f> ColorCorrector::extractColorSamples(const cv::Mat& colorCard) {
    std::vector<cv::Vec3f> colors;
    
    if (colorCard.empty()) {
        return colors;
    }
    
    // 将色卡分割成6x4的网格 (24个色块)
    int rows = 10;
    int cols = 8;
    int blockHeight = colorCard.rows / rows;
    int blockWidth = colorCard.cols / cols;
    
    for (int r = 0; r < rows; ++r) {
        for (int c = 0; c < cols; ++c) {
            // 计算每个色块的中心区域 (避免边缘噪声)
            int startX = c * blockWidth + blockWidth / 4;
            int startY = r * blockHeight + blockHeight / 4;
            int endX = (c + 1) * blockWidth - blockWidth / 4;
            int endY = (r + 1) * blockHeight - blockHeight / 4;
            
            cv::Rect roi(startX, startY, endX - startX, endY - startY);
            cv::Mat colorBlock = colorCard(roi);
            
            // 计算该色块的平均颜色
            cv::Scalar meanColor = cv::mean(colorBlock);
            // OpenCV使用BGR顺序，但我们需要保存为RGB顺序用于颜色校正
            colors.push_back(cv::Vec3f(meanColor[2], meanColor[1], meanColor[0])); // BGR -> RGB
        }
    }
    
    return colors;
}

bool ColorCorrector::calculateColorCorrectionMatrix(const cv::Mat& referenceCard, const cv::Mat& inputCard) {
    // 提取参考色卡和输入色卡的颜色样本
    std::vector<cv::Vec3f> refColors = extractColorSamples(referenceCard);
    std::vector<cv::Vec3f> inputColors = extractColorSamples(inputCard);
    
    if (refColors.size() != inputColors.size() || refColors.empty()) {
        std::cout << "[错误] 颜色样本数量不匹配或为空" << std::endl;
        return false;
    }
    
    // 输出前几个颜色样本进行调试
    std::cout << "[调试] 提取的颜色样本数量: " << refColors.size() << std::endl;
    for (int i = 0; i < std::min(80, (int)refColors.size()); ++i) {
        std::cout << "[调试] 参考色卡第" << i << "个颜色 (RGB): (" 
                  << refColors[i][0] << ", " << refColors[i][1] << ", " << refColors[i][2] << ")" << std::endl;
        std::cout << "[调试] 输入色卡第" << i << "个颜色 (RGB): (" 
                  << inputColors[i][0] << ", " << inputColors[i][1] << ", " << inputColors[i][2] << ")" << std::endl;
    }
    
    // 使用最小二乘法计算3x3色彩校正矩阵
    // 构建方程组 Ax = b，其中A是输入颜色矩阵，b是参考颜色矩阵
    
    int numSamples = static_cast<int>(inputColors.size());
    cv::Mat A(numSamples, 3, CV_32F);
    cv::Mat bR(numSamples, 1, CV_32F);
    cv::Mat bG(numSamples, 1, CV_32F);
    cv::Mat bB(numSamples, 1, CV_32F);
    
    // 填充矩阵 (输入和参考颜色都使用RGB顺序)
    for (int i = 0; i < numSamples; ++i) {
        A.at<float>(i, 0) = inputColors[i][0]; // R
        A.at<float>(i, 1) = inputColors[i][1]; // G
        A.at<float>(i, 2) = inputColors[i][2]; // B
        
        bR.at<float>(i, 0) = refColors[i][0]; // R
        bG.at<float>(i, 0) = refColors[i][1]; // G
        bB.at<float>(i, 0) = refColors[i][2]; // B
    }
    
    // 使用最小二乘法求解每个颜色通道的变换系数
    cv::Mat xR, xG, xB;
    cv::solve(A, bR, xR, cv::DECOMP_SVD);
    cv::solve(A, bG, xG, cv::DECOMP_SVD);
    cv::solve(A, bB, xB, cv::DECOMP_SVD);
    
    // 构建3x3色彩校正矩阵
    colorCorrectionMatrix = cv::Mat::zeros(3, 3, CV_32F);
    
    // 填充矩阵 (注意OpenCV使用BGR顺序，所以需要调整通道顺序)
    colorCorrectionMatrix.at<float>(0, 0) = xB.at<float>(2, 0); // R->B
    colorCorrectionMatrix.at<float>(0, 1) = xB.at<float>(1, 0); // G->B
    colorCorrectionMatrix.at<float>(0, 2) = xB.at<float>(0, 0); // B->B
    
    colorCorrectionMatrix.at<float>(1, 0) = xG.at<float>(2, 0); // R->G
    colorCorrectionMatrix.at<float>(1, 1) = xG.at<float>(1, 0); // G->G
    colorCorrectionMatrix.at<float>(1, 2) = xG.at<float>(0, 0); // B->G
    
    colorCorrectionMatrix.at<float>(2, 0) = xR.at<float>(2, 0); // R->R
    colorCorrectionMatrix.at<float>(2, 1) = xR.at<float>(1, 0); // G->R
    colorCorrectionMatrix.at<float>(2, 2) = xR.at<float>(0, 0); // B->R
    
    isCalibrated = true;
    
    std::cout << "[信息] 色彩校正矩阵计算完成:" << std::endl;
    std::cout << colorCorrectionMatrix << std::endl;
    
    return true;
}

cv::Mat ColorCorrector::applyColorCorrection(const cv::Mat& image) {
    if (!isCalibrated) {
        std::cout << "[警告] 未进行校准，返回原始图像" << std::endl;
        return image.clone();
    }
    
    cv::Mat result;
    image.convertTo(result, CV_32F);
    
    // 应用色彩校正矩阵到每个像素
    std::vector<cv::Mat> channels;
    cv::split(result, channels);
    
    cv::Mat correctedChannels[3];
    for (int i = 0; i < 3; ++i) {
        correctedChannels[i] = cv::Mat::zeros(result.size(), CV_32F);
        for (int j = 0; j < 3; ++j) {
            correctedChannels[i] += colorCorrectionMatrix.at<float>(i, j) * channels[j];
        }
        // 限制值在0-255范围内
        cv::threshold(correctedChannels[i], correctedChannels[i], 255, 255, cv::THRESH_TRUNC);
        cv::threshold(correctedChannels[i], correctedChannels[i], 0, 0, cv::THRESH_TOZERO);
    }
    
    // 合并通道
    std::vector<cv::Mat> correctedChannelVec = {correctedChannels[0], correctedChannels[1], correctedChannels[2]};
    cv::merge(correctedChannelVec, result);
    
    // 转换回8位
    result.convertTo(result, CV_8U);
    
    return result;
}

bool ColorCorrector::saveParametersToJson(const std::string& filename) {
    if (!isCalibrated) {
        std::cout << "[错误] 未进行校准，无法保存参数" << std::endl;
        return false;
    }
    
    json j;
    
    // 保存色彩校正矩阵
    j["color_correction_matrix"] = json::array();
    for (int i = 0; i < 3; ++i) {
        json row = json::array();
        for (int k = 0; k < 3; ++k) {
            row.push_back(colorCorrectionMatrix.at<float>(i, k));
        }
        j["color_correction_matrix"].push_back(row);
    }
    
    // 保存校准状态
    j["is_calibrated"] = isCalibrated;
    
    // 保存时间戳
    auto now = std::chrono::system_clock::now();
    auto time_t = std::chrono::system_clock::to_time_t(now);
    j["timestamp"] = std::ctime(&time_t);
    
    // 写入文件
    try {
        std::ofstream file(filename);
        file << j.dump(4);
        file.close();
        std::cout << "[信息] 校正参数已保存到: " << filename << std::endl;
        return true;
    } catch (const std::exception& e) {
        std::cout << "[错误] 保存JSON文件失败: " << e.what() << std::endl;
        return false;
    }
}

bool ColorCorrector::loadParametersFromJson(const std::string& filename) {
    try {
        std::ifstream file(filename);
        if (!file.is_open()) {
            std::cout << "[错误] 无法打开JSON文件: " << filename << std::endl;
            return false;
        }
        
        json j;
        file >> j;
        file.close();
        
        // 加载色彩校正矩阵
        if (j.contains("color_correction_matrix")) {
            colorCorrectionMatrix = cv::Mat::zeros(3, 3, CV_32F);
            auto matrix = j["color_correction_matrix"];
            for (int i = 0; i < 3; ++i) {
                for (int k = 0; k < 3; ++k) {
                    colorCorrectionMatrix.at<float>(i, k) = matrix[i][k];
                }
            }
        }
        
        // 加载校准状态
        if (j.contains("is_calibrated")) {
            isCalibrated = j["is_calibrated"];
        }
        
        std::cout << "[信息] 校正参数已从文件加载: " << filename << std::endl;
        std::cout << "[信息] 色彩校正矩阵:" << std::endl;
        std::cout << colorCorrectionMatrix << std::endl;
        
        return true;
    } catch (const std::exception& e) {
        std::cout << "[错误] 加载JSON文件失败: " << e.what() << std::endl;
        return false;
    }
}

void ColorCorrector::showComparison(const cv::Mat& original, const cv::Mat& corrected) {
    // 创建并排显示的图像
    cv::Mat comparison;
    cv::hconcat(original, corrected, comparison);
    
    // 在图像上添加标签
    cv::putText(comparison, "Original", cv::Point(30, 30), 
                cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 255, 0), 2);
    cv::putText(comparison, "Corrected", cv::Point(original.cols + 30, 30), 
                cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 255, 0), 2);
    
    cv::imshow("Color Correction Comparison", comparison);
}
